package com.macro.mall.common.aop;

import com.google.common.util.concurrent.RateLimiter;
import com.macro.mall.common.annotation.RateLimiterDIY;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: TuDou
 * @create: 2024-10-25 16:45
 *
 * 高频接口限流
 *
 *     @RateLimiterDIY(value = 10, durationInSeconds = 60) //限制为每分钟 10 次请求
 **/
@Aspect
@Component
public class RateLimiterAspect {

    private final ConcurrentHashMap<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();

    @Pointcut("@annotation(rateLimiterAnnotation)")
    public void rateLimiterPointcut(RateLimiterDIY rateLimiterAnnotation) {
    }

    @Around("rateLimiterPointcut(rateLimiterAnnotation)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimiterDIY rateLimiterAnnotation) throws Throwable {
        int permits = rateLimiterAnnotation.value();
        int durationInSeconds = rateLimiterAnnotation.durationInSeconds();

        // 使用方法签名作为 RateLimiter 的 key
        String key = joinPoint.getSignature().toLongString();
        RateLimiter rateLimiter = rateLimiters.computeIfAbsent(key, k -> RateLimiter.create((double) permits / durationInSeconds));

        // 尝试获取令牌，如果获取到则执行方法，否则抛出异常
        if (rateLimiter.tryAcquire()) {
            return joinPoint.proceed();
        } else {
            throw new RuntimeException("Rate limit exceeded.");
        }
    }
}

